﻿using System;
using System.Collections.Generic;
using MLSharp.Utilities;

namespace MLSharp
{
	/// <summary>
	/// Provides utility methods to partition a data
	/// set into test and training data for the 
	/// purpose of performing cross-fold validation.
	/// </summary>
	public static class KFold
	{
		/// <summary>
		/// Breaks the specified data set up into the specified
		/// number of folds and returns the partitions.
		/// </summary>
		/// <param name="dataSet">The data set to partition.</param>
		/// <param name="numFolds">The number of folds to create.</param>
		/// <returns>An array of partitions.</returns>
		public static Partition[] GetPartitions(IDataSet dataSet, int numFolds)
		{
			//Clone the data set before partitioning it.
			dataSet = dataSet.Clone();

			//Validate parameters
			if (numFolds <= 0 || numFolds > dataSet.Instances.Count)
			{
				throw (new ArgumentOutOfRangeException("numFolds", numFolds,
				                                       "Value must be greater than zero and less than the size of the data set."));
			}

			//Step 1: Break the data set into N folds
			List<Instance>[] folds = new List<Instance>[numFolds];
			for (int i=0; i < folds.Length; i++)
			{
				folds[i] = new List<Instance>();
			}

			Sample sample = new Sample(dataSet.Instances);
			for (int i = 0; i < dataSet.Instances.Count; i++)
			{
				folds[i%numFolds].Add(sample.SampleWithoutReplacement());
			}

			//Step 2: Create N partitions using each of the folds as a validation
			//set exactly once
			Partition[] partitions = new Partition[numFolds];
			for (int i = 0; i < partitions.Length; i++)
			{
				partitions[i] = new Partition(dataSet.Attributes);

				List<Instance> training = new List<Instance>();

				for (int j=0; j < folds.Length; j++)
				{
					if (i != j)
						training.AddRange(folds[j]);
					else
						partitions[i].Validation.Instances.AddRange(folds[j]);
				}

				partitions[i].Training.Instances.AddRange(training);

				//Copy the target index from the original data set.
				partitions[i].Training.TargetAttributeIndex =
					partitions[i].Validation.TargetAttributeIndex = dataSet.TargetAttributeIndex;
			}

			//Step 3: Return the result
			return partitions;
		}
	}
}
